import json
import threading
from enum import Enum
import opuslib
import pyaudio
import requests
import websocket
import configs
import logging
import colorlog
from logging.handlers import TimedRotatingFileHandler
import sys
from pynput import keyboard

_formatter = colorlog.ColoredFormatter(
    fmt='%(log_color)s%(asctime)s %(levelname)s: %(message)s',
    log_colors={
        'DEBUG': 'white',
        'INFO': 'green',
        'WARNING': 'yellow',
        'ERROR': 'red',
        'CRITICAL': 'bold_red',
    },
)

_file_formatter = logging.Formatter('%(asctime)s %(levelname)s: %(message)s')

_file_handler = TimedRotatingFileHandler("runlog.log", when="midnight", interval=1, backupCount=30, encoding='utf-8')
_file_handler.setLevel(logging.DEBUG)
_file_handler.setFormatter(_file_formatter)

_console_handler = logging.StreamHandler(stream=sys.stdout)
_console_handler.setLevel(logging.DEBUG)
_console_handler.setFormatter(_formatter)

_logger = logging.getLogger('logger')
_logger.setLevel('DEBUG')
_logger.addHandler(_console_handler)
_logger.addHandler(_file_handler)

headers = {
    "Authorization": f"Bearer {configs.TOKEN}",
    "Protocol-Version": "1",
    "Device-Id": configs.DEVICE_ID,
    "Client-Id": configs.CLIENT_ID
}


class StatusEnum(Enum):
    IDLE = 0
    CONNECTING = 1
    CONNECTED = 2
    HANDSHAKE = 3
    LISTENING = 4
    RECORD = 5
    PLAY = 6


sample_rate = 24000
channels = 1
chunk = int(60 * sample_rate / 1000)

status = StatusEnum.IDLE
session_id = ''

ws = None


def check_ota_version():
    body = {
        "flash_size": 16777216,
        "minimum_free_heap_size": 8318916,
        "mac_address": f"{configs.DEVICE_ID}",
        "chip_model_name": "esp32s3",
        "chip_info": {
            "model": 9,
            "cores": 2,
            "revision": 2,
            "features": 18
        },
        "application": {
            "name": "xiaozhi",
            "version": "1.0.0"
        },
        "partition_table": [],
        "ota": {
            "label": "factory"
        },
        "board": {
            "type": "xiaozhi-python",
            "ip": "192.168.124.38",
            "mac": f"{configs.DEVICE_ID}"
        }
    }
    headers = {
        'Device-Id': f"{configs.DEVICE_ID}",
        'Content-Type': 'application/json'
    }
    response = requests.post(configs.URL_OTA, headers=headers, data=json.dumps(body))


def on_open(client: websocket.WebSocket):
    global status
    status = StatusEnum.CONNECTED
    _logger.debug('websocket connect successfully.')


def on_close(client: websocket.WebSocket, status_code, message):
    global status
    status = StatusEnum.IDLE
    _logger.debug(f'websocket closed, status_code:{status_code}, message:{message}')


def on_error(client: websocket.WebSocket, error):
    _logger.error(f'websocket error:{error}')


def on_reconnect(client: websocket.WebSocket):
    _logger.debug(f'websocket reconnect')


def on_message(client: websocket.WebSocket, message):
    global chunk, sample_rate, channels, session_id, status
    if isinstance(message, bytes):
        data = decoder.decode(message, chunk)
        audio_stream_output.write(data)
    else:
        try:
            _logger.debug(f'recv msg:{message}')
            response = json.loads(message)
            if response['type'] == 'hello':
                session_id = response['session_id']
                audio_params = response['audio_params']
                sample_rate = audio_params['sample_rate']
                channels = audio_params['channels']
                chunk = int(audio_params['frame_duration'] * audio_params['sample_rate'] / 1000)
                status = StatusEnum.LISTENING
        except Exception as e:
            _logger.critical(e)


def send_message(data):
    global ws
    raw = json.dumps(data)
    _logger.debug(f'send msg:{raw}')
    ws.send(raw)


def send_opus(data):
    global ws
    ws.send(data, opcode=websocket.ABNF.OPCODE_BINARY)


def keyboard_on_press_callback(key: keyboard.KeyCode):
    global status, ws

    if key != keyboard.Key.space:
        return

    if status == StatusEnum.IDLE:
        check_ota_version()
        ws = websocket.WebSocketApp(
            configs.URL_WEB,
            header=headers,
            on_open=on_open,
            on_close=on_close,
            on_message=on_message,
            on_error=on_error,
            on_reconnect=on_reconnect,
        )
        t = threading.Thread(target=ws.run_forever)
        t.start()
        status = StatusEnum.CONNECTING
        return

    if status == StatusEnum.CONNECTED:
        msg = {
            "type": "hello",
            "version": 1,
            "transport": "websocket",
            "audio_params": {
                "format": "opus",
                "sample_rate": configs.SAMPLE_RATE,
                "channels": configs.CHANNELS,
                "frame_duration": configs.FRAME_DURATION,
            }
        }
        send_message(msg)
        status = StatusEnum.HANDSHAKE
        return

    # 通知服务器开始监听
    if status == StatusEnum.LISTENING:
        msg = {
            "session_id": session_id,
            "type": "listen",
            "state": "start",
            "mode": "manual"
        }
        send_message(msg)
        status = StatusEnum.RECORD
        _logger.debug('请说话...')
        return

    # 开始说话
    if status == StatusEnum.RECORD:
        data = audio_stream_input.read(chunk)
        data = encoder.encode(data, chunk)
        send_opus(data)
        return


def keyboard_on_release_callback(key: keyboard.KeyCode):
    global status
    if key != keyboard.Key.space:
        return
    # 通知服务器停止监听
    msg = {
        "session_id": session_id,
        "type": "listen",
        "state": "stop",
        "mode": "manual"
    }
    send_message(msg)
    status = StatusEnum.LISTENING
    _logger.debug('说话结束！')


# 打开音频输入流
audio_stream_input = pyaudio.PyAudio().open(
    channels=channels,
    rate=sample_rate,
    frames_per_buffer=chunk,
    format=pyaudio.paInt16,
    input=True,
)

# 打开音频输出流
audio_stream_output = pyaudio.PyAudio().open(
    channels=channels,
    rate=sample_rate,
    frames_per_buffer=chunk,
    format=pyaudio.paInt16,
    output=True,
)

# 创建音频编码器
encoder = opuslib.Encoder(sample_rate, channels, opuslib.APPLICATION_AUDIO)

# 创建音频解码器
decoder = opuslib.Decoder(sample_rate, channels)

# 监听键盘按键
_logger.debug('长按空格键说话...')
listener = keyboard.Listener(on_press=keyboard_on_press_callback, on_release=keyboard_on_release_callback)
listener.start()
listener.join()
